#include <linux/init.h>
#include <linux/module.h>
#include <linux/vmalloc.h>
#include <linux/blkdev.h>
#include <linux/hdreg.h>
#include <linux/blk-mq.h>

static struct gendisk *mydisk;
unsigned int major;
void *disk_addr;

#define BLKSIZE (8 * 1024 * 1024)

int mydisk_open(struct block_device *bdev, fmode_t mode)
{
    printk("enter %s\n", __func__);
    return 0;
}

void mydisk_release(struct gendisk *disk, fmode_t mode)
{
    printk("enter %s\n", __func__);
}

int mydisk_getgeo(struct block_device *bdev, struct hd_geometry *hd)
{
    hd->heads = 4;
    hd->cylinders = 16;
    hd->sectors = BLKSIZE / hd->heads / hd->cylinders / 512;
    return 0;
}

struct block_device_operations mem_blk_fops = 
{
    .open = mydisk_open,
    .release = mydisk_release,
    .getgeo = mydisk_getgeo,
};

static struct blk_mq_tag_set set;

blk_status_t mydisk_queue_rq(struct blk_mq_hw_ctx *hctx,
                             const struct blk_mq_queue_data *bd)
{
    struct request *rq = bd->rq;
    struct req_iterator iter;   // 迭代器
    struct bio_vec bvec;

    loff_t pos = blk_rq_pos(rq) * 512;

    blk_mq_start_request(rq);

    rq_for_each_segment(bvec, rq, iter) 
    {
        void *b_buf = page_address(bvec.bv_page) + bvec.bv_offset;
        unsigned int b_len = bvec.bv_len;

        if(b_len > BLKSIZE)
            b_len = BLKSIZE;

        if(rq_data_dir(rq)) 
            memcpy(disk_addr + pos, b_buf, b_len);
        else
            memcpy(b_buf, disk_addr + pos, b_len);

        pos += b_len;
    }

    blk_mq_end_request(rq, BLK_STS_OK);

    return BLK_STS_OK;
}

static struct blk_mq_ops myqops = 
{
    .queue_rq = mydisk_queue_rq,
};

int __init mydisk_init(void)
{
    mydisk = alloc_disk(4);

    major = register_blkdev(0, "mydisk");
    mydisk->major = major;
    mydisk->first_minor = 0;
    strcpy(mydisk->disk_name, "mydisk");
    set_capacity(mydisk, BLKSIZE >> 9);
    mydisk->fops = &mem_blk_fops;

    mydisk->queue = blk_mq_init_sq_queue(&set, &myqops, 16, BLK_MQ_F_SHOULD_MERGE);

    disk_addr = vmalloc(BLKSIZE);

    add_disk(mydisk);
    return 0;
}

void __exit mydisk_exit(void)
{
    del_gendisk(mydisk);
    vfree(disk_addr);
    blk_cleanup_queue(mydisk->queue);
    unregister_blkdev(major, "mydisk");

    put_disk(mydisk);
}

module_init(mydisk_init);
module_exit(mydisk_exit);

MODULE_LICENSE("GPL");